# -*- coding: utf-8 -*-
"""
技术分析模块 - 基于您提供的指标公式进行技术分析
"""
import pandas as pd
import numpy as np
import talib as ta
from typing import Dict, Tuple
import logging
from config import TECHNICAL_INDICATORS

logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)

class TechnicalAnalyzer:
    """技术分析器 - 基于您的指标公式"""
    
    def __init__(self):
        self.config = TECHNICAL_INDICATORS
    
    def calculate_all_indicators(self, data: pd.DataFrame) -> pd.DataFrame:
        """计算所有技术指标"""
        if data.empty or len(data) < 60:
            logger.warning("数据不足，无法计算技术指标")
            return data
        
        df = data.copy()
        
        # 基础数据检查
        required_columns = ['open', 'high', 'low', 'close', 'volume']
        for col in required_columns:
            if col not in df.columns:
                logger.error(f"缺少必要列: {col}")
                return data
        
        # 1. 计算您的买卖线指标
        df = self._calculate_buy_sell_lines(df)
        
        # 2. 计算指导线
        df = self._calculate_guide_lines(df)
        
        # 3. 计算多重EMA系统
        df = self._calculate_ema_system(df)
        
        # 4. 计算买卖信号
        df = self._calculate_signals(df)
        
        # 5. 计算额外的技术指标
        df = self._calculate_additional_indicators(df)
        
        return df
    
    def _calculate_buy_sell_lines(self, df: pd.DataFrame) -> pd.DataFrame:
        """计算买线和卖线"""
        try:
            # 买线 = EMA(C,2)
            df['买线'] = ta.EMA(df['close'].values, timeperiod=2)
            
            # 卖线 = EMA(SLOPE(C,21)*20+C,42)
            # 首先计算21周期的斜率
            slope_period = self.config['slope_period']
            slope_multiplier = self.config['slope_multiplier']
            
            # 计算斜率 (简化版本: 线性回归斜率)
            slopes = []
            close_values = df['close'].values
            
            for i in range(len(close_values)):
                if i < slope_period - 1:
                    slopes.append(np.nan)
                else:
                    # 计算过去N天的线性回归斜率
                    y = close_values[i-slope_period+1:i+1]
                    x = np.arange(len(y))
                    if len(y) > 1:
                        slope = np.polyfit(x, y, 1)[0]
                        slopes.append(slope)
                    else:
                        slopes.append(0)
            
            df['slope_21'] = slopes
            
            # 卖线 = EMA(SLOPE(C,21)*20+C,42)
            df['slope_adjusted'] = df['slope_21'] * slope_multiplier + df['close']
            df['卖线'] = ta.EMA(df['slope_adjusted'].values, timeperiod=42)
            
            logger.info("买卖线计算完成")
            return df
            
        except Exception as e:
            logger.error(f"买卖线计算失败: {e}")
            return df
    
    def _calculate_guide_lines(self, df: pd.DataFrame) -> pd.DataFrame:
        """计算指导线和界线"""
        try:
            # 指导 = EMA((EMA(CLOSE,4)+EMA(CLOSE,6)+EMA(CLOSE,12)+EMA(CLOSE,24))/4,2)
            ema4 = ta.EMA(df['close'].values, timeperiod=4)
            ema6 = ta.EMA(df['close'].values, timeperiod=6)
            ema12 = ta.EMA(df['close'].values, timeperiod=12)
            ema24 = ta.EMA(df['close'].values, timeperiod=24)
            
            avg_ema = (ema4 + ema6 + ema12 + ema24) / 4
            df['指导'] = ta.EMA(avg_ema, timeperiod=2)
            
            # 界 = MA(CLOSE,27)
            df['界'] = ta.SMA(df['close'].values, timeperiod=27)
            
            logger.info("指导线计算完成")
            return df
            
        except Exception as e:
            logger.error(f"指导线计算失败: {e}")
            return df
    
    def _calculate_ema_system(self, df: pd.DataFrame) -> pd.DataFrame:
        """计算多重EMA系统"""
        try:
            # P1: EMA(CLOSE, 5)
            df['P1'] = ta.EMA(df['close'].values, timeperiod=5)
            
            # P2: EMA(CLOSE, 8) 
            df['P2'] = ta.EMA(df['close'].values, timeperiod=8)
            
            # P3: EMA(CLOSE, 55)
            df['P3'] = ta.EMA(df['close'].values, timeperiod=55)
            
            # 计算价格与EMA的关系
            df['price_vs_p1'] = df['close'] - df['P1']
            df['price_vs_p2'] = df['close'] - df['P2'] 
            df['price_vs_p3'] = df['close'] - df['P3']
            
            # 颜色信号
            df['p1_signal'] = np.where(df['close'] <= df['P1'], 'GREEN', 'RED')
            df['p2_signal'] = np.where(df['close'] <= df['P2'], 'GREEN', 'RED')
            df['p3_signal'] = np.where(df['close'] <= df['P3'], 'GREEN', 'RED')
            
            logger.info("EMA系统计算完成")
            return df
            
        except Exception as e:
            logger.error(f"EMA系统计算失败: {e}")
            return df
    
    def _calculate_signals(self, df: pd.DataFrame) -> pd.DataFrame:
        """计算买卖信号"""
        try:
            # 买入信号: CROSS(买线, 卖线) - 买线上穿卖线
            df['buy_signal'] = False
            df['sell_signal'] = False
            
            if '买线' in df.columns and '卖线' in df.columns:
                # 计算交叉信号
                buy_cross = (df['买线'] > df['卖线']) & (df['买线'].shift(1) <= df['卖线'].shift(1))
                sell_cross = (df['卖线'] > df['买线']) & (df['卖线'].shift(1) <= df['买线'].shift(1))
                
                df['buy_signal'] = buy_cross
                df['sell_signal'] = sell_cross
            
            # 强度信号
            df['signal_strength'] = 0
            
            # 多头排列加分
            multi_bull = (df['P1'] > df['P2']) & (df['P2'] > df['P3']) & (df['close'] > df['P1'])
            df.loc[multi_bull, 'signal_strength'] += 3
            
            # 价格突破指导线
            guide_break = df['close'] > df['指导']
            df.loc[guide_break, 'signal_strength'] += 2
            
            # 成交量配合
            vol_ma5 = df['volume'].rolling(5).mean()
            vol_ma20 = df['volume'].rolling(20).mean()
            vol_surge = vol_ma5 > vol_ma20 * 1.2
            df.loc[vol_surge, 'signal_strength'] += 1
            
            logger.info("交易信号计算完成")
            return df
            
        except Exception as e:
            logger.error(f"信号计算失败: {e}")
            return df
    
    def _calculate_additional_indicators(self, df: pd.DataFrame) -> pd.DataFrame:
        """计算额外的技术指标"""
        try:
            # RSI
            df['RSI'] = ta.RSI(df['close'].values, timeperiod=14)
            
            # MACD
            df['MACD'], df['MACD_signal'], df['MACD_hist'] = ta.MACD(df['close'].values)
            
            # 布林带
            df['BB_upper'], df['BB_middle'], df['BB_lower'] = ta.BBANDS(df['close'].values)
            
            # KDJ
            df['K'], df['D'] = ta.STOCH(df['high'].values, df['low'].values, df['close'].values)
            df['J'] = 3 * df['K'] - 2 * df['D']
            
            # ATR (真实波动范围)
            df['ATR'] = ta.ATR(df['high'].values, df['low'].values, df['close'].values, timeperiod=14)
            
            # 成交量指标
            df['volume_ma5'] = df['volume'].rolling(5).mean()
            df['volume_ma20'] = df['volume'].rolling(20).mean()
            df['volume_ratio'] = df['volume'] / df['volume_ma20']
            
            # 价格变化率
            df['price_change_1d'] = df['close'].pct_change()
            df['price_change_5d'] = df['close'].pct_change(5)
            df['price_change_20d'] = df['close'].pct_change(20)
            
            logger.info("额外技术指标计算完成")
            return df
            
        except Exception as e:
            logger.error(f"额外指标计算失败: {e}")
            return df
    
    def get_current_signals(self, df: pd.DataFrame) -> Dict:
        """获取当前的交易信号"""
        if df.empty:
            return {}
        
        latest = df.iloc[-1]
        
        signals = {
            'buy_signal': bool(latest.get('buy_signal', False)),
            'sell_signal': bool(latest.get('sell_signal', False)),
            'signal_strength': latest.get('signal_strength', 0),
            'trend_status': self._get_trend_status(latest),
            'risk_level': self._get_risk_level(latest),
            'recommendation': self._get_recommendation(latest)
        }
        
        return signals
    
    def _get_trend_status(self, latest_data: pd.Series) -> str:
        """获取趋势状态"""
        try:
            if latest_data['close'] > latest_data['P1'] > latest_data['P2'] > latest_data['P3']:
                return "强势上涨"
            elif latest_data['close'] > latest_data['P2'] > latest_data['P3']:
                return "上涨趋势"
            elif latest_data['close'] > latest_data['P3']:
                return "弱势上涨"
            elif latest_data['close'] < latest_data['P3']:
                return "下跌趋势"
            else:
                return "震荡"
        except:
            return "未知"
    
    def _get_risk_level(self, latest_data: pd.Series) -> str:
        """获取风险等级"""
        try:
            atr = latest_data.get('ATR', 0)
            price = latest_data.get('close', 0)
            
            if price > 0:
                volatility = atr / price
                if volatility > 0.05:
                    return "高风险"
                elif volatility > 0.03:
                    return "中风险"
                else:
                    return "低风险"
            return "未知"
        except:
            return "未知"
    
    def _get_recommendation(self, latest_data: pd.Series) -> str:
        """获取投资建议"""
        try:
            buy_signal = latest_data.get('buy_signal', False)
            sell_signal = latest_data.get('sell_signal', False)
            strength = latest_data.get('signal_strength', 0)
            rsi = latest_data.get('RSI', 50)
            
            if buy_signal and strength >= 4:
                return "强烈买入"
            elif buy_signal and strength >= 2:
                return "买入"
            elif sell_signal:
                return "卖出"
            elif rsi > 70:
                return "超买，谨慎"
            elif rsi < 30:
                return "超卖，关注"
            else:
                return "持有观望"
        except:
            return "持有观望"

if __name__ == "__main__":
    # 测试技术分析
    analyzer = TechnicalAnalyzer()
    
    # 创建测试数据
    test_data = pd.DataFrame({
        'open': np.random.randn(100).cumsum() + 100,
        'high': np.random.randn(100).cumsum() + 102,
        'low': np.random.randn(100).cumsum() + 98,
        'close': np.random.randn(100).cumsum() + 100,
        'volume': np.random.randint(1000, 10000, 100)
    })
    
    result = analyzer.calculate_all_indicators(test_data)
    signals = analyzer.get_current_signals(result)
    
    print("技术指标列:", result.columns.tolist())
    print("当前信号:", signals)
